﻿#include "dwt.hpp"
#include "common.hpp"
//#include "filter_types.hpp"

namespace cuda
{

/*
 one round 5/3 wavelet decomposition using lifting
*/

__device__
int wavelet_analysis( float2 * data, int n,   // original data
				      float2 * approximation, // approximation coefficients
				      float2 * detail )       // detail coefficients
{
	//const float2 two = {2, 2};
	int m, i;
	float2 * x, * d;
	
	m = n/2;
	n = m * 2;  /* n must be even */

	x = (float2*) malloc( sizeof(float2) * (n+3) );
	d = (float2*) malloc( sizeof(float2) * (m+1) );
// 	x = (float2*) realloc( x, sizeof(float2) * (n+3) );
// 	d = (float2*) realloc( d, sizeof(float2) * (m+1) );

	for ( i = 0; i < n; i++ ) /* boundary extension, CB ABCDEF E */
	{
//		x[i+2] = data[i];
		x[i+2].x = data[i].x;
		x[i+2].y = data[i].y;
	}
/*	x[0] = data[2];*/
	x[0].x = data[2].x;
	x[0].y = data[2].y;
/*	x[1] = data[1];*/
	x[1].x = data[1].x;
	x[1].y = data[1].y;
/*	x[n+2] = data[n-2];*/
	x[n+2].x = data[n-2].x;
	x[n+2].y = data[n-2].y;
	/* detail (odd) */
	for ( i = 0; i <= m; i++ )
	{
/*		d[i] = x[i*2+1] - (x[i*2]+x[i*2+2])/2;*/
		d[i].x = x[i*2+1].x - (x[i*2].x+x[i*2+2].x)/2;
		d[i].y = x[i*2+1].y - (x[i*2].y+x[i*2+2].y)/2;
	}
	for ( i = 0; i < m; i++ )
	{
/*		detail[i] = d[i+1];*/
		detail[i].x = d[i+1].x;
		detail[i].y = d[i+1].y;
	}
	/* approximation (even) */
	for ( i = 0; i < m; i++ )
	{
//		approximation[i] = x[2*i+2] + (d[i+1] + d[i]+ two)/4;
		approximation[i].x = x[2*i+2].x + ( d[i+1].x+d[i].x )/4;
		approximation[i].y = x[2*i+2].y + ( d[i+1].y+d[i].y )/4;
	}
	free( x );
	free( d );

	return( m );
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/*
 Wavelet synthesis
*/
__device__
int wavelet_synthesis( float2 * approximation, // approximation coefficients
				       float2 * detail,        // detail coefficients
					   int n,               // n is length of both detail and approximation
				       float2 * data )         // synthesised data
{
	float2 * s, * d, * x;
	int i, m;

	m = n * 2;
	s = (float2 *) malloc( sizeof(float2) * (n+1) );
	d = (float2 *) malloc( sizeof(float2) * (n+2) );
	x = (float2 *) malloc( sizeof(float2) * (m+1) );

	/* boundary extension */
	for ( i = 0; i < n; i++ ) 
	{
/*		s[i] = approximation[i];*/
		s[i].x = approximation[i].x;
		s[i].y = approximation[i].y;
	}
	s[n] = approximation[n-1]; /* right extension, ...EF F */
	for ( i = 0; i < n; i++ ) 
	{
/* 		d[i+1] = detail[i];*/
		d[i+1].x = detail[i].x;
		d[i+1].y = detail[i].y;
	}
/*	d[0] = detail[0]; */ /* left extension, A ABC ... */
	d[0].x = detail[0].x;
	d[0].y = detail[0].y;
/*	d[n+1] = detail[n-2]; *//* right extension, ... DEF E */
	d[n+1].x = detail[n-2].x;
	d[n+1].y = detail[n-2].y;
	/* lifting */
	for ( i = 0; i < n+1; i++ )
	{
//		x[i*2] = s[i]-(d[i]+d[i+1]+two)/4;
		x[i*2].x = s[i].x-(d[i].x+d[i+1].x+2)/4;
		x[i*2].y = s[i].y-(d[i].y+d[i+1].y+2)/4;
	}
	for ( i = 0; i < n; i++ )
	{
/*		x[i*2+1] = d[i+1] + (x[i*2]+x[i*2+2])/2;*/
		x[i*2+1].x = d[i+1].x + (x[i*2].x+x[i*2+2].x)/2;
		x[i*2+1].y = d[i+1].y + (x[i*2].y+x[i*2+2].y)/2;
	}
	/* copy x to data */
	for ( i = 0; i < m; i++ )
	{
/*		data[i] = x[i];*/
		data[i].x = x[i].x;
		data[i].y = x[i].y;
	}

	free( x );
	free( s );
	free( d );

	//return( m );
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/*
 multi-level decomposition for 2D signal
*/
__global__
void multilevel_decomposition_2D( float2 * image,			// original signal
								 int width, int height,		// number of pointers, should be multiple of 2
								 float2 * wavelet_data,		// decomposited data
								 int levels )				// decomposition levels  
{
	int i, x, y, ww, hh;
	int rtv;
	float2 * l, * h, * data;

	ww = ( width>>levels  )<<levels; // guartee nn to be multiple of 2
	hh = ( height>>levels )<<levels;

	for ( y = 0; y < height; y++ )
	{
		for ( x = 0; x < width; x++ )
		{
/*			wavelet_data[y*width+x] = image[y*width+x];*/
			wavelet_data[y*width+x].x = image[y*width+x].x;
			wavelet_data[y*width+x].y = image[y*width+x].y;
		}
	}

	for ( i = 0; i < levels; i++ )
	{
		for ( y = 0; y < hh; y++ )
		{
			l = (float2*)malloc(sizeof(float2)*ww);
			h = (float2*)malloc(sizeof(float2)*ww);
			data = (float2*)malloc(sizeof(float2)*ww);
			for ( x = 0; x < ww; x++ ) 
			{
/*				data[x] = wavelet_data[y*width+x];*/
				data[x].x = wavelet_data[ y*width+x ].x;
				data[x].y = wavelet_data[ y*width+x ].y;
			}
			rtv = wavelet_analysis( data, ww, l, h );
			for ( x = 0; x < ww/2; x++ )
			{
/*				wavelet_data[y*width+x] = l[x];*/
				wavelet_data[y*width+x].x = l[x].x;
				wavelet_data[y*width+x].y = l[x].y;
			}
			for ( x = 0; x < ww/2; x++ ) 
			{
/*				wavelet_data[y*width+x+ww/2] = h[x];*/
				wavelet_data[y*width+x+ww/2].x = h[x].x;
				wavelet_data[y*width+x+ww/2].y = h[x].y;
			}
			free(l);
			free(h);
			free(data);
		}
		for ( x = 0; x < ww; x++ )
		{
			l = (float2*)malloc(sizeof(float2)*hh);
			h = (float2*)malloc(sizeof(float2)*hh);
			data = (float2*)malloc(sizeof(float2)*hh);
			for ( y = 0; y < hh; y++ ) 
			{
/*				data[y] = wavelet_data[ y*width+x ];*/
				data[y].x = wavelet_data[ y*width+x ].x;
				data[y].y = wavelet_data[ y*width+x ].y;
			}
			rtv = wavelet_analysis( data, hh, l, h );
			for ( y = 0; y < hh/2; y++ ) 
			{
/*				wavelet_data[y*width+x] = l[y];*/
				wavelet_data[y*width+x].x = l[y].x;
				wavelet_data[y*width+x].y = l[y].y;
			}
			for ( y = 0; y < hh/2; y++ ) 
			{
/*				wavelet_data[(y+hh/2)*width+x] = h[y];*/
				wavelet_data[(y+hh/2)*width+x].x = h[y].x;
				wavelet_data[(y+hh/2)*width+x].y = h[y].y;
			}
			free(l);
			free(h);
			free(data);
		}
		ww = ww/2;
		hh = hh/2;
	}
	//__syncthreads();
	//return( levels ); // return decomposition levels
}

/*
 multi-level reconstruction for 2D signal
*/
__global__
void multilevel_reconstruction_2D( float2 * wavelet_data,  // wavelet coefficients
								  int width, int height,  // number of pointers, should be multiple of 2
								  float2 * image,         // reconstructed data
								  int levels )            // reconstruction levels
{
	int i, x, y, ww, hh;
	int rtv;
	float2 *l, *h, *data;

	ww = width>>levels;  // guartee nn to be multiple of 2
	hh = height>>levels;

	for ( y = 0; y < height; y++ )
	{
		for ( x = 0; x < width; x++ )
		{
/*			image[y*width+x] = wavelet_data[y*width+x];*/
			image[y*width+x].x = wavelet_data[y*width+x].x;
			image[y*width+x].y = wavelet_data[y*width+x].y;
		}
	}
	for ( i = 0; i < levels; i++ )
	{

		for ( x = 0; x < ww*2; x++ )
		{
			l = (float2*)malloc(sizeof(float2)*hh);
			h = (float2*)malloc(sizeof(float2)*hh);
			data = (float2*)malloc(sizeof(float2)*2*hh);
			for ( y = 0; y < hh; y++ ) 
			{
// 				l[y] = image[ y*width+x ];
// 				h[y] = image[ (y+hh)*width+x ];
				l[y].x = image[ y*width+x ].x;
				h[y].x = image[ (y+hh)*width+x ].x;
				l[y].y = image[ y*width+x ].y;
				h[y].y = image[ (y+hh)*width+x ].y;
			}
			rtv = wavelet_synthesis( l, h, hh, data );
			for ( y = 0; y < hh*2; y++ ) 
			{
/*				image[y*width+x] = data[y];*/
				image[y*width+x].x = data[y].x;
				image[y*width+x].y = data[y].y;
			}
			free(l);
			free(h);
			free(data);
		}

		for ( y = 0; y < hh*2; y++ )
		{
			l = (float2*)malloc(sizeof(float2)*ww);
			h = (float2*)malloc(sizeof(float2)*ww);
			data = (float2*)malloc(sizeof(float2)*2*ww);
			for ( x = 0; x < ww; x++ )
			{
// 				l[x] = image[ y*width+x ];
//  			h[x] = image[ y*width+x+ww ];
				l[x].x = image[ y*width+x ].x;
				h[x].x = image[ y*width+x+ww ].x;
				l[x].y = image[ y*width+x ].y;
				h[x].y = image[ y*width+x+ww ].y;
			}
			rtv = wavelet_synthesis( l, h, ww, data );
			for ( x = 0; x < ww*2; x++ ) 
			{
/*				image[y*width+x] = data[x];*/
				image[y*width+x].x = data[x].x;
				image[y*width+x].y = data[x].y;
			}
			free(l);
			free(h);
			free(data);
		}
		ww = ww*2;
		hh = hh*2;
	}
	//syncthreads();
	//return( levels ); // return decomposition levels
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// <summary>	Wavelet Transform </summary>
/// <remarks>	Quan, 6/18/2012. </remarks>
////////////////////////////////////////////////////////////////////////////////////////////////////
#define N (512*512)
#define THREADS_PER_BLOCK 512
void dwt(
	float2* src,	
	float2* dst,	
	int nRows,
	int nCols,
	//int nChannels,
	int nLevels,
	int flag)
{
	/// <summary>	Prepare for Wavelet Transform </summary>


	/// <summary>	Perform Wavelet Transform </summary>
	switch(flag)
	{
	/// <summary>	Forward Wavelet Transform </summary>
	case DWT_FORWARD:		
		multilevel_decomposition_2D <<<1, 1>>>	
			(src,
			nCols, nRows,
			dst,
			nLevels);
		break;
	
	/// <summary>	Inverse Wavelet Transform </summary>
 	case DWT_INVERSE:		
		multilevel_reconstruction_2D <<<1, 1>>>	
			(src,
			nCols, nRows,
			dst,
			nLevels);
 		break;

// 	default:
// 		break;
	}
}
}
